import argparse
import igraph
import ipaddress
import random
import csv
import os.path

from SimulationConfiguration import Conf
from SSFNetDML import SSFNetDML

class AS:

    THRESHOLD_FOR_COMPLEX_NETWORK = 2

    def __init__(self, asn, name, is_tier_1, is_preferred=False, file_num=0):
        self.asn = asn
        self.name = name
        self.is_tier_1 = is_tier_1
        self.arin_rank = 0
        self.out_rank = 0
        self.costumers = set()
        self.peers = set()
        self.providers = set()
        self.main_provider = None
        self.unknown_relations = set()
        self.prefixes = set()
        self.is_preferred = is_preferred
        self.file_num = file_num

    def __eq__(self,other):
        if type(other) == type(self):
            return self.asn == other.asn and self.name == other.name
        if type(other) == int:
            return self.asn == other
        return False
    def __hash__(self):
        return self.asn.__hash__()
    def __str__(self):
        return "{0} {1} arin rank = {2} out rank = {3}".format(self.asn, self.name, self.arin_rank, self.out_rank)
    def __repr__(self):
        return str(self)

class AS_link:
    def __init__(self,as0,as1):
        self.as0 = as0
        self.as1 = as1
    def __eq__(self,other):
        if type(other) == type(self):
            return (self.as0 == other.as0 and self.as1 == other.as1) or self.as0 == other.as1 and self.as1 == other.as0
        if type(other) == AS:
            return self.as0 == other or self.as1 == other
        return False
    def __hash__(self):
        return self.as0.__hash__() * self.as1.__hash__()
        
def getAllLinks(ASes, as_links_file):
    all_links = []
    all_arin_links = []
    ASes_dict = dict((x, x) for x in ASes)
    with open(as_links_file,'r') as links_file:
        lines = links_file.readlines()
        for line in lines:
            all_links.append(list(map(lambda x: int(x), line.split())))
            link = all_links[-1]
            
            if link[0] in ASes_dict and link[1] in ASes_dict:
                AS0 = ASes_dict.get(link[0])
                AS1 = ASes_dict.get(link[1])
                AS0.arin_rank += 1
                AS1.arin_rank += 1
                all_arin_links.append(link)
            elif link[0] in ASes_dict:
                AS0 = ASes_dict.get(link[0])
                AS0.out_rank += 1
            elif link[1] in ASes_dict:
                AS0 = ASes_dict.get(link[1])
                AS0.out_rank += 1
            
    return all_arin_links

def setASRelationships(ASes, as_relationships_file):
    with open(as_relationships_file,'r') as relationships_file:
        lines = relationships_file.readlines()
    for line in lines:
        as0, as1, rel = line.split()
        as0 = int(as0)
        as1 = int(as1)
        if as0 in ASes and as1 in ASes:
            AS0 = ASes[ASes.index(as0)]
            AS1 = ASes[ASes.index(as1)]
            if rel == 'p2c':
                AS0.costumers.add(AS1)
                AS1.providers.add(AS0)
            elif rel == 'c2p':
                AS1.costumers.add(AS0)
                AS0.providers.add(AS1)
            elif rel == 'p2p':
                AS0.peers.add(AS1)
                AS1.peers.add(AS0)
            else:
                AS0.unknown_relations.add(AS1)
                AS1.unknown_relations.add(AS0)
    # Select main providers
    for AS in ASes:
        if AS.providers:
            AS.main_provider = random.choice(list(AS.providers))
            AS.providers.remove(AS.main_provider)

def setPrefixOwnership(mini_ASes, as_origin_file):
    with open(as_origin_file,'r') as origin_file:
        lines = origin_file.readlines()
    for line in lines:
        lineSplit = line.split()
        asn = int(lineSplit[1])
        if asn in mini_ASes:
            index = mini_ASes.index(asn)
            mini_ASes[index].prefixes.add(ipaddress.ip_network(lineSplit[0]))
    
def getAllAses(ASes_files, ranking_file):
    ASes = []
    index = 0
    with open(ranking_file, 'r') as ranking_file:
        for row in csv.reader(ranking_file):
            # Hack: all asns in one row
            ranking = list(map(lambda st: int(st), row))
            tier_1 = ranking[:Conf.TIER_1_THRESHOLD]
            break
    for ases_file in ASes_files:
        with open(ases_file,'r') as as_file:
            lines = as_file.readlines()
            for line in lines:
                as_record = line.split()
                asn = int(as_record[1])
                ASes.append(AS(asn, as_record[0], asn in tier_1, file_num=index))
        index += 1
    return ASes

def prefer_ASes(ASes, preferrred_ASes=[]):
    preferrred_ASes = set(preferrred_ASes)
    ASes = set(ASes)

    intersection = list(set(ASes).intersection(preferrred_ASes))
    difference = ASes.difference(intersection)

    intersection_length = len(intersection)

    if intersection_length >= Conf.MAX_NEIGHBOURS:
        return intersection
    return list(intersection) + list(difference)[:Conf.MAX_NEIGHBOURS - intersection_length]

def getMiniTopology(seed_AS, ASes, links, max_ASes, preferred_ASes):
    ASes_dict = dict((x, x) for x in ASes)
    mini_ASes = [ASes_dict[seed_AS]]
    links_dict = {}
    for link in links:
        if link[0] in links_dict:
            links_dict[ASes_dict[link[0]]].append(ASes_dict[link[1]])
        else:
            links_dict[ASes_dict[link[0]]] = [ASes_dict[link[1]]]
            
        if link[1] in links_dict:
            links_dict[ASes_dict[link[1]]].append(ASes_dict[link[0]])
        else:
            links_dict[ASes_dict[link[1]]] = [ASes_dict[link[0]]]
    to_add = []
    links_to_add = []
    mini_links = []
    while len(mini_ASes) < max_ASes:
        for AS_ in mini_ASes:
            to_add = to_add + prefer_ASes(links_dict[AS_], preferred_ASes)
            links_to_add = links_to_add + list(map(lambda x: AS_link(AS_,x),links_dict[AS_]))
        to_add = list(set(to_add))
        links_to_add = list(set(links_to_add))
        number_to_add = max(0, min(len(to_add), max_ASes - len(mini_ASes)))

        # Make sure we don't remove the preferred ASes
        to_add.sort(key=lambda AS: AS.is_preferred, reverse=True)
        mini_ASes = mini_ASes + to_add[:number_to_add]

        links_to_add = list(filter(lambda link: link.as0 in mini_ASes and link.as1 in mini_ASes,links_to_add))
        mini_links = mini_links + links_to_add
        to_add = []
        links_to_add = []
    return list(set(mini_ASes)), list(set(mini_links))

def rank_color(rank):
    if rank == 1:
        return 'green'
    if rank == 2:
        return 'blue'
    return 'red'
def marked_color(asn, file_num, is_preferred):
    if asn == '3356' or asn == '174':
        return 'orange'
    if asn == '378' or asn == '20965':
        return 'purple'
    if is_preferred:
        return 'red'

    return {0: 'gray', 1: 'pink', 2: 'yellow', 3: 'white'}[file_num]
    
    
def plot_topology(mini_ASes, mini_links ):
    g = igraph.Graph()
    for AS_node in mini_ASes:
        g.add_vertex(name = str(AS_node.asn), AS_name = AS_node.name, arin_rank = AS_node.arin_rank, out_rank = AS_node.out_rank, label = str(AS_node.asn) + '\n' + AS_node.name, is_preferred = AS_node.is_preferred, file_num = AS_node.file_num)
    for link in mini_links:
        g.add_edge(str(link.as0.asn), str(link.as1.asn),width=2)
    visual_style = {}
    visual_style["vertex_size"] = 70
    visual_style["vertex_color"] = list(map(lambda v: marked_color(v['name'], v['file_num'], v['is_preferred']),g.vs))
    #visual_style["layout"] = igraph.layout
    visual_style["bbox"] = (2000, 2000)
    visual_style["margin"] = 30
    layout = g.layout("kamada_kawai")
    visual_style["layout"] = layout
    igraph.plot(g, "arin.png",**visual_style)

def getPreferredASes(preferred_ASes_file):
    ASes = []
    if not preferred_ASes_file:
        return ASes
    AS_reader = csv.reader(preferred_ASes_file, delimiter=' ')
    for row in AS_reader:
        asn = int(row[0])
        ASes.append(AS(asn, row[1], False, True))
    return ASes

def main(ases_files, as_links_file, as_relationships_file, as_origin_file, ranking_file, seed_AS, max_ASes, output_dml, plot, preferred_ASes_file):
    ASes = set(getAllAses(ases_files, ranking_file))
    preferred_ASes = set(getPreferredASes(preferred_ASes_file))
    ASes = list(preferred_ASes.union(ASes))

    links = getAllLinks(ASes, as_links_file)

    mini_ASes, mini_links = getMiniTopology(int(seed_AS), ASes, links, int(max_ASes), preferred_ASes)
    setASRelationships(mini_ASes, as_relationships_file)
    setPrefixOwnership(mini_ASes, as_origin_file)
    if plot:
        plot_topology(mini_ASes, mini_links )
    with SSFNetDML(output_dml) as DML:
        DML.set_ASes(mini_ASes)
        DML.set_links(mini_links)
    

if (__name__ == "__main__"):
    parser = argparse.ArgumentParser()
    parser.add_argument('ases_files', nargs='+', type=str, help='List of ASes files')
    parser.add_argument("as_links_file",)
    parser.add_argument("as_relationships_file",)
    parser.add_argument("as_origin_file",)
    parser.add_argument("ranking_file",)
    parser.add_argument("seed_AS",)
    parser.add_argument("max_ASes",)
    parser.add_argument("output_dml",)
    parser.add_argument("-p", "--plot", help="increase output verbosity",
                    action="store_true")
    parser.add_argument('--preferred_ASes_file', nargs='?', type=argparse.FileType('r'), help='ASNs as the first word of each line')
    args = parser.parse_args()
    main(**args.__dict__)